//
// Created by wangwanlong on 2024/11/4.
//

#ifndef MKAHYPAR_SIMPLE_REFINER_H
#define MKAHYPAR_SIMPLE_REFINER_H

#include "mkahypar/partition/refinement/i_refiner.h"
#include "mkahypar/partition/refinement/refiner_base.h"
#include "mkahypar/partition/refinement/policies/block_penalty_policy.h"
#include "mkahypar/partition/refinement/policies/edge_penalty_policy.h"
#include "mkahypar/partition/refinement/policies/hop_penalty_policy.h"
#include "mkahypar/partition/refinement/policies/stop_policy.h"
#include "mkahypar/datastructure/sparse_set.h"


namespace mkahypar {
  template<class BlockPenaltyPolicy = NormalBlockPenalty,
          class EdgePenaltyPolicy = NormalEdgePenaltyPolicy,
          class HopPenaltyPolicy = NormalHopPenaltyPolicy,
          class StoppingPolicy = Mandatory>
  class SimpleRefiner : public IRefiner,
                        public RefinerBase {
  public:
    SimpleRefiner(Hypergraph &hypergraph, TargetGraph &topo) :
            RefinerBase(hypergraph, topo), _stopping_policy(),
            _active_nodes(hypergraph.initialNumNodes()),
            _best_partitions(hypergraph.initialNumNodes()) {

    }

    //拿起一个节点
    void pickNode(HypernodeID hn, PartitionID from, PartitionID to) {
      if (to != Hypergraph::kInvalidPartition || from == Hypergraph::kInvalidPartition)
        ERROREXIT("error!");
      //资源量更正
      _topo._fpga[from]._used_res -= _hg.hypernode(hn).weight();
      //更改节点分区
      _hg.hypernode(hn).part_id = to;
      for (const HyperedgeID &he: _hg.incidentEdges(hn)) {
        HypernodeID source = _hg.hyperedge(he).source();
        HyperedgeWeight he_weight = _hg.hyperedge(he).weight();
        //更新pin_in_part
        bool no_pins_left_in_source_part = _hg.decrementPinCountInPart(he, from);
        bool only_one_pin_in_to_part = _hg.incrementPinCountInPart(he, to);
        //更新target_part_route_count
        if (hn != source) {
          _hg.decrementRouteCountInPart(he, from, hn);
          _hg.incrementRouteCountInPart(he, to, hn);
        }
      }
    }

    //寻找最好的分区
    void findBestMove(HypernodeID v, Metrics &best_gain, PartitionID &best_part) {
      PartitionID from = _hg.partID(v);
      for (PartitionID id = 0; id < _hg._k; ++id) {
        if (id == from)
          continue;
        Metrics delta_gain;
        evaluateNodeMoveGain(v, from, id, delta_gain);
        if (delta_gain.obj > best_gain.obj) {
          best_part = id;
          best_gain = delta_gain;
        }
      }
    }


    //评估节点移动收益,这个函数不考虑from和to都是非法的情况
    void evaluateNodeMoveGain(HypernodeID v, PartitionID from, PartitionID to, Metrics &delta_gain) {

      //先计算资源的变化
      delta_gain.block_penalty += BlockPenaltyPolicy::computeMove(_hg.nodeWeight(v),
                                                                  _topo._fpga[from]._used_res,
                                                                  _topo._fpga[from]._bound_res);
      //计算资源的变化量
      delta_gain.block_penalty += BlockPenaltyPolicy::computeAdd(_hg.nodeWeight(v),
                                                                 _topo._fpga[to]._used_res,
                                                                 _topo._fpga[to]._bound_res);
      GainType from_line_delta = 0, to_line_delta = 0;
      //计算路由和超载
      for (auto &he: _hg.incidentEdges(v)) {
        auto &edge = _hg.hyperedge(he);
        auto he_weight = _hg.edgeWeight(he);
        auto source = _hg.hyperedge(he).source();
        HypernodeID source_id = edge.source();
        bool no_pins_left_in_source_part = _hg.pinCountInPart(he, from) == 1;
        bool only_one_pin_in_to_part = _hg.pinCountInPart(he, to) == 0;
        if (no_pins_left_in_source_part && !only_one_pin_in_to_part) {
          if (_hg.pinCountInPart(he, to) + 1 == _hg.edgeSize(he)) {
            //从边界边 -> 内部边, from和to的超载都不需要计算
            from_line_delta -= he_weight;
            to_line_delta -= he_weight;
          } else {
            //边界边 -> 边界边，只不过不需要计算from的超载量
            from_line_delta -= he_weight;
          }
        } else if (!no_pins_left_in_source_part && only_one_pin_in_to_part) {
          if (_hg.pinCountInPart(he, from) == _hg.edgeSize(he)) {
            //从内部边 -> 边界边, from和to都需要计算超载量
            from_line_delta += he_weight;
            to_line_delta += he_weight;
          } else {
            //从边界边->边界边，只不过需要计算to的超载量
            to_line_delta += he_weight;
          }
        }else if(no_pins_left_in_source_part && only_one_pin_in_to_part){
          //从边界边->边界边，减少from的互联量，增加to的互联量
          from_line_delta -= he_weight;
          to_line_delta += he_weight;
        }
        //计算路由变化
        if(v == source){
          for (PartitionID pid = 0; pid < _hg._k; ++pid) {
            GainType count = _hg.routeCountInPart(he, pid);
            if (count > 0) {
              delta_gain.hop_penalty += HopPenaltyPolicy::compute(he_weight, _topo._distance[from][pid]) * count;
              delta_gain.hop_penalty -= HopPenaltyPolicy::compute(he_weight, _topo._distance[to][pid]) * count;
            }
          }
        }
        else{
          PartitionID  source_part = _hg.partID(source);
          GainType count = _hg.tarRouteCount(he, v);
          delta_gain.hop_penalty += HopPenaltyPolicy::compute(he_weight, _topo._distance[source_part][from])*count;
          delta_gain.hop_penalty -= HopPenaltyPolicy::compute(he_weight, _topo._distance[source_part][to])*count;
        }
      }
      //计算From互联的变化量
      delta_gain.edge_penalty += EdgePenaltyPolicy::computeDelta(_topo._fpga[from]._used_line,
                                                                 _topo._fpga[from]._used_line + from_line_delta,
                                                                 _topo._fpga[from]._bound_line);
      //计算to分区的互联变化量
      delta_gain.edge_penalty += EdgePenaltyPolicy::computeDelta(_topo._fpga[to]._used_line,
                                                                 _topo._fpga[to]._used_line + to_line_delta,
                                                                 _topo._fpga[to]._bound_line);
      delta_gain.updateObj();
    }

    //验证数据结构
    void check(Metrics &current_obj) {
      //检查超边的数据结构
      int num_partitions = _hg._k;
      for (const HyperedgeID &he: _hg.edges()) {
        int edge_size = 0;
        auto source = _hg.hyperedge(he).source();
        std::vector<int> pin_in_part_size(num_partitions);
        std::vector<int> part_route_count(num_partitions);
        int64_t connect = 0;
        for (auto hn: _hg.pins(he)) {
          ++edge_size;
          pin_in_part_size[_hg.partID(hn)] += 1;
          connect |= 1LL << _hg.partID(hn);
          if (hn != source)
            part_route_count[_hg.partID(hn)] += _hg.tarRouteCount(he, hn);
        }
        //检查联通集
        if (connect != _hg.hyperedge(he).connectivity._n) {
          LOG << "connectivity  error!";
          LOG << he << connect << _hg.hyperedge(he).connectivity._n;
          ERROREXIT("");
        }
        //检查edge size
        if (edge_size != _hg.edgeSize(he)) {
          LOG << "edge size error!";
          LOG << he << edge_size << _hg.edgeSize(he);
          ERROREXIT("");
        }
        for (int id = 0; id < num_partitions; ++id) {
          //检查pin size
          if (pin_in_part_size[id] != _hg.pinCountInPart(he, id)) {
            LOG << "pin in part size error!";
            LOG << he << pin_in_part_size[id] << _hg.pinCountInPart(he, id);
            ERROREXIT("");
          }
          //检查part route count
          if (part_route_count[id] != _hg.routeCountInPart(he, id)) {
            LOG << "route count in part size error!";
            LOG << he << id << part_route_count[id] << _hg.routeCountInPart(he, id);
            ERROREXIT("");
          }
        }
      }
      //检查超点的数据结构
      for (const HypernodeID &hn: _hg.nodes()) {
        int num_incident_cuts = 0;
        for (const auto &he: _hg.hypernode(hn).incidentNets()) {
          if (_hg.connectivity(he) > 1)
            ++num_incident_cuts;
        }
        if (num_incident_cuts != _hg.hypernode(hn).num_incident_cut_hes) {
          LOG << "node cut error!";
          LOG << hn << num_incident_cuts << _hg.hypernode(hn).num_incident_cut_hes;
          ERROREXIT("error!");
        }
      }

      //检查互联量
      std::vector<int> line(num_partitions);
      for (const HyperedgeID &he: _hg.edges()) {
        auto he_weight = _hg.hyperedge(he).weight();
        auto connect = _hg.hyperedge(he).connectivity._n;
        if (_hg.connectivity(he) > 1) {
          for (int id = 0; id < num_partitions; ++id) {
            bool t1 = _hg.pinCountInPart(he,id) > 0;
            bool t2 = (connect >> id) & 1;
            if(t1 != t2){
              LOG<<t1<<t2;
//              LOG<<id<<_hg.pinCountInPart(he,id)<<_hg.connectivity(he).;
              ERROREXIT("error!");
            }
            if(_hg.pinCountInPart(he,id) > 0) {
              line[id] += he_weight;
            }
          }
        }
      }
      for (int id = 0; id < num_partitions; ++id) {
        LOG<<"pid:"<<id<<_topo._fpga[id]._used_line;
        if(line[id] != _topo._fpga[id]._used_line){
          LOG<<"line error!";
          LOG<<id<<line[id]<<_topo._fpga[id]._used_line;

          LOG<<"12 size"<<_hg.edgeSize(12);

          LOG<<"14 adj edge:";
          for(auto ee : _hg.incidentEdges(14))
            LLOG<<ee;
          LOG<<"15 adj edge:";
          for(auto ee : _hg.incidentEdges(15))
            LLOG<<ee;
          _hg.printGraph();
//          _topo.prinInfo();

          ERROREXIT("error!");
        }
      }
      //验证obj 计算是否正确
      NodeWeight s;
      Metrics obj;
      for (PartitionID id = 0; id < _topo._fpga_num; ++id) {
        //计算block约束
        obj.block_penalty += BlockPenaltyPolicy::exceed(_topo._fpga[id]._used_res,
                                                        _topo._fpga[id]._bound_res);
        //计算超载约束
        obj.edge_penalty += EdgePenaltyPolicy::exceed(_topo._fpga[id]._used_line,
                                                      _topo._fpga[id]._bound_line);
      }
      //计算路由约束
      for (const auto &he: _hg.edges()) {
        HypernodeID source = _hg.hyperedge(he).source();
        HyperedgeWeight he_weight = _hg.edgeWeight(he);
        for (PartitionID pid = 0; pid < _hg._k; ++pid) {
          GainType count = _hg.routeCountInPart(he, pid);
          if (count > 0) {
            GainType distance = _topo._distance[_hg.partID(source)][pid];
//            LOG<<"he:"<<he<<"he_weight:"<<he_weight<<"distance:"<<_topo._distance[_hg.partID(source)][pid]<<"count:"<<count;
            obj.hop_penalty += HopPenaltyPolicy::compute(he_weight, distance) * count;
          }
        }
      }
      obj.updateObj();
      if (obj.obj != current_obj.obj) {
        LOG << "obj error!";
        LOG << "true obj:";
        obj.print();
        LOG << "current obj:";
        current_obj.print();
        ERROREXIT("");
      }

    }


    //初始解之后或者还原之后执行全局搜索
    void globalSearch(std::vector<HypernodeID> &refinement_nodes,Metrics &best_obj){

    }


  protected:
    //局部搜索
    bool refineImpl(std::vector<HypernodeID> &refinement_nodes,
                    Metrics &best_obj) override {
//      LOG<<"before refine check:";
//      check(best_obj);

      _active_nodes.clear();
      //加入新的点到活跃点集合
      for (auto v: refinement_nodes) {
        if (_hg.isBorderNode(v))
          _active_nodes.add(v);
      }
      //记录最优解的index
      int min_obj_index = -1;
      _performed_moves.clear();

      //重置停止策略
      _stopping_policy.resetStatistics();
      int touched_hns_since_last_improvement = 0;
      Metrics current_obj = best_obj;
      Metrics local_best = best_obj;
      const double beta = log(_hg._num_hypernodes);

      while (_active_nodes.size() > 0) {
        std::vector<HypernodeID> nodes;
        for (auto v: _active_nodes) {
          if (_hg.isBorderNode(v))
            nodes.emplace_back(v);
        }
        bool isEnd = false;
        Randomize::getInstance().shuffleVector(nodes, nodes.size());
        for (auto v: nodes) {
          if (_stopping_policy.searchShouldStop(touched_hns_since_last_improvement,
                                                beta, best_obj.obj, current_obj.obj)) {
            isEnd = true;
            break;
          }
          Metrics best_gain;
          PartitionID best_part = Hypergraph::kInvalidPartition;
          best_gain.obj = std::numeric_limits<GainType>::min();
          //搜索最佳移动分区
          findBestMove(v, best_gain, best_part);

          if (best_part == Hypergraph::kInvalidPartition)
            ERROREXIT("error!");

          //更新停止策略
          ++touched_hns_since_last_improvement;
          _stopping_policy.updateStatistics(best_gain.obj);

          if (best_gain.obj >= 0) {//只执行能够改善目标函数的解
            //更新目标函数并且执行移动
            current_obj -= best_gain;
            PartitionID old_part = _hg.partID(v);
            changeNodePart(v, _hg.partID(v), best_part);
            //保存移动
            _performed_moves.emplace_back(RollbackInfo{v, old_part, _hg.partID(v)});



//            LOG << "try move node:" << v << old_part << "to" << best_part;
//            LOG<<"after move check";
//            check(current_obj);
            //保存举报最优解
            if (current_obj.obj < best_obj.obj) {
              //重置停止策略
              _stopping_policy.resetStatistics();
              touched_hns_since_last_improvement = 0;
              //记录最优解的索引,回滚的时候需要
              min_obj_index = _performed_moves.size() - 1;
              best_obj = current_obj;
            }
            //更新活跃点
          }
        }
        if (isEnd)
          break;
      }
      //回滚最优解
      int last_index = _performed_moves.size() - 1;
      while (last_index != min_obj_index) {
        const HypernodeID hn = _performed_moves[last_index].hn;
        const PartitionID from_part = _performed_moves[last_index].to_part;
        const PartitionID to_part = _performed_moves[last_index].from_part;
        changeNodePart(hn, from_part, to_part);
        --last_index;
      }
      //check

//      LOG << "final check";
//      check(best_obj);

      return true;

    }

    //非增量计算目标函数值
    void computeObjImpl(Metrics &obj) override {
      NodeWeight s;

      for (PartitionID id = 0; id < _topo._fpga_num; ++id) {
        //计算block约束
        obj.block_penalty += BlockPenaltyPolicy::exceed(_topo._fpga[id]._used_res,
                                                        _topo._fpga[id]._bound_res);
        //计算超载约束
        obj.edge_penalty += EdgePenaltyPolicy::exceed(_topo._fpga[id]._used_line,
                                                      _topo._fpga[id]._bound_line);
      }
      //计算路由约束
      for (const auto &he: _hg.edges()) {
        HypernodeID source = _hg.hyperedge(he).source();
        for (PartitionID pid = 0; pid < _hg._k; ++pid) {
          GainType count = _hg.routeCountInPart(he, pid);
          if (count > 0) {
            GainType distance = _topo._distance[_hg.partID(source)][pid];
            obj.hop_penalty += HopPenaltyPolicy::compute(_hg.hyperedge(he).weight(),
                                                         distance) * count;
          }
        }
      }
      obj.updateObj();
    }


    StoppingPolicy _stopping_policy;
    ds::SparseSet<HypernodeID> _active_nodes;//活跃节点集合
    std::vector<HypernodeID> _best_partitions;

  };


}

#endif //MKAHYPAR_SIMPLE_REFINER_H
